import pinocchio as pin
import numpy as np
from pinocchio.utils import rand, zero
import meshcat

# 引入src中所有的函数
import sys
from os.path import dirname, join

sys.path.append(join(dirname(dirname(__file__)), "src"))
sys.path.append(dirname(dirname(__file__))) # 把文件根路径添加到sys.path中，这样在根路径下的src文件夹才能被import
from src import *
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# 获取URDF模型路径
URDF_FILE_PATH = filePath(URDF_NAME = 'ur5_robot', ROBOT = True)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# 构建robot对象，并得到viz
robot, viz = robot_viz(URDF_FILE_PATH)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
# 前向运动学
init_config = Init_Config(ROBOT=True)

idx = robot.index(init_config.joint_names[1])

q1 = np.array([0, 0, 0, -1, 0, 0])

# 这里是给定各个关节角和特定关节index，然后求这个特定关节的姿态
# 该方法会有返回，可以写成placement=robot.placement(q0, idx)，但是这样再直接修改placement就会改变robot姿态
# 使用Placement同时就会改变robot姿态
robot.placement(q1, idx)
placement = robot.data.oMi[idx].copy()
print(placement)
print('type of placement is ', type(placement))
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

############################################################
# 创建一个小球
sphere_name = "world/sphere"
sphere_radius = 0.1

viz.viewer["world/sphere"].set_object(
    meshcat.geometry.Sphere(sphere_radius),
    meshcat.geometry.MeshLambertMaterial(color=0x6622dd,reflectivity=0.4)
)

# 设置小球的位置和姿态
position = np.array([0.5, 0.1, 0.2])
quaternion = [1.0, 0.0, 0.0, 0.0]  # w, x, y, z
rotation_matrix = pin.Quaternion(quaternion[0],quaternion[1],quaternion[2],quaternion[3]).matrix()
viz.viewer["world/sphere"].set_transform(
    pin.SE3(rotation_matrix, position).homogeneous
)

# 显示机器人模型，并修改机器人姿态
# 需要指出的是，meshcat的viz.display这里传入的是6维ndarray，不是7维向量
# 另外不要弄混np.array([])和np.ndarray([])

robot.initViewer(loadModel=True)
viz.display(q1)
q2 = np.array([0.5, -0.5, -0.5, 0, 0, 0])
viz.display(q2)